#!/usr/bin/env python3
"""
Arm Point Cloud Trimming Script

This script provides two methods to trim arm points from point cloud data based on URDF and joint states:
1. Mesh-based trimming: Uses robot mesh geometry dynamically generated from URDF
2. Cylindrical spine-based trimming: Uses cylindrical approximation of arm links with torus joints

Usage:
    trimArmPcd(unifiedData, param, urdf_path, pcd_keys)
    
Parameters:
    - 'm100': Use mesh-based trimming with 100mm buffer
    - 'c100': Use cylindrical trimming with 100mm radius
"""

import numpy as np
import logging
import traceback
from typing import Dict, List, Optional, Tuple, Union, Sequence
from pathlib import Path
from time_it import timeit
import os
import time
from scalePcd import scalePcd

# Robot kinematics imports
import yourdfpy#type: ignore
from scipy.spatial.transform import Rotation as R
from scipy.spatial.distance import cdist
import trimesh
import trimesh.creation

# Try to import PyTorch3D for vectorized operations
import torch
import torch.nn.functional as F

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Global debugging variables for spine and torus control
spine_mask = True        # Enable/disable spine (cylinder) masking
spine_vis = False       # Enable/disable spine visualization

# Global variables for mesh simplification control
MESH_SIMPLIFY_TARGET_FACES = 5      # Target number of faces for mesh simplification
MESH_SIMPLIFY_MAX_FACES = 15        # Maximum number of faces allowed
MESH_VISUALIZATION_ENABLED = True    # Enable/disable mesh visualization
MESH_SURFACE_POINTS_PER_FACE = 10    # Number of points per face when visualizing mesh surface

# Global variables for debugging and visualization
DEBUG_PRINT_MESH_INFO = False          # Print mesh information during processing
DEBUG_PRINT_SPINE_INFO = False        # Print spine point information
DEBUG_PRINT_FILTERING_INFO = False    # Print filtering statistics

class ArmPointCloudTrimmer:
    """
    Trims arm points from point cloud data using URDF and joint states.
    Uses torus shapes between cylindrical segments for better performance.
    """
    
    def __init__(self, urdf_path: str):
        """
        Initialize the trimmer with URDF only.
        
        Args:
            urdf_path: Path to the URDF file
        """
        self.urdf_path = urdf_path
        
        # Load URDF
        self.robot = yourdfpy.URDF.load(urdf_path, build_scene_graph=True, load_meshes=False)
        
        # Also load a version with meshes for mesh-based trimming
        try:
            self.robot_with_meshes = yourdfpy.URDF.load(urdf_path, build_scene_graph=True, load_meshes=True)
        except Exception as e:
            logger.warning(f"Could not load robot with meshes: {e}")
            self.robot_with_meshes = None
        
        # Define arm joint names based on the URDF structure
        self.arm_joint_names = [
            'panda_joint1', 'panda_joint2', 'panda_joint3', 'panda_joint4', 
            'panda_joint5', 'panda_joint6', 'panda_joint7'
        ]
        
        # Define arm link names that contain the physical arm geometry
        self.arm_link_names = [
            'panda_link0', 'panda_link1', 'panda_link2', 'panda_link3', 
            'panda_link4', 'panda_link5', 'panda_link6', 'panda_link7', 
            'panda_hand', 'panda_leftfinger', 'panda_rightfinger'
        ]
        
        # Define left and right arm poses in the workspace
        # Based on the config file robot_pose: [0, -0.65, 0.75, 0.707, 0, 0, 0.707]
        # For dual arm setup, we'll offset them slightly
        arm_separation_offset = 0.10  # Additional separation distance in meters (adjustable)
        self.left_arm_base_pose = np.array([-0.3 - arm_separation_offset, -0.65, 0.75])  # Left arm offset (moved left by 0.15m)
        self.right_arm_base_pose = np.array([0.3 + arm_separation_offset, -0.65, 0.75])  # Right arm offset (moved right by 0.15m)
        
        # Create rotation matrices for 90-degree counter-clockwise rotation (viewed from above)
        # Rotation around Z-axis by 90 degrees (pi/2 radians)
        self.rotation_90_ccw = np.array([
            [0, -1, 0],
            [1, 0, 0],
            [0, 0, 1]
        ])
        
        # Create full transformation matrices for both arms
        self.left_arm_transform = np.eye(4)
        self.left_arm_transform[:3, :3] = self.rotation_90_ccw
        self.left_arm_transform[:3, 3] = self.left_arm_base_pose
        
        self.right_arm_transform = np.eye(4)
        self.right_arm_transform[:3, :3] = self.rotation_90_ccw
        self.right_arm_transform[:3, 3] = self.right_arm_base_pose

    def trim_with_mesh(self, pcd_array: np.ndarray, joint_states: np.ndarray, buffer_mm: float = 100.0) -> np.ndarray:
        """
        Trim points using mesh-based collision detection with robot meshes.
        
        Args:
            pcd_array: Point cloud array of shape (N, 6) with [x, y, z, r, g, b]
            joint_states: Array of shape (16,) containing [left_arm_7_joints, left_gripper_1, right_arm_7_joints, right_gripper_1]
            buffer_mm: Buffer distance in millimeters
            
        Returns:
            Filtered point cloud array
        """
        global MESH_VISUALIZATION_ENABLED, DEBUG_PRINT_MESH_INFO
        
        if pcd_array.shape[0] == 0:
            return pcd_array
            
        points_xyz = pcd_array[:, :3]
        buffer_m = buffer_mm / 1000.0  # Convert to meters
        
        # Split joint states for left and right arms (excluding gripper values)
        left_joint_states = joint_states[:7]  # First 7 values are left arm joints
        right_joint_states = joint_states[8:15]  # Values 8-14 are right arm joints (skip gripper at index 7)
        
        # Get meshes for both arms
        left_arm_meshes = self._get_arm_meshes(left_joint_states, self.left_arm_base_pose, is_left_arm=True)
        right_arm_meshes = self._get_arm_meshes(right_joint_states, self.right_arm_base_pose, is_left_arm=False)
        
        # If no meshes could be loaded, fall back to cylindrical method
        if len(left_arm_meshes) == 0 and len(right_arm_meshes) == 0:
            logger.warning("No meshes could be loaded, falling back to cylindrical method")
            return self.trim_with_cylinder(pcd_array, joint_states, buffer_mm)
        
        # Generate mesh visualization points if enabled
        mesh_vis_points = np.empty((0, 6))
        if MESH_VISUALIZATION_ENABLED:
            if DEBUG_PRINT_MESH_INFO:
                print(f"DEBUG: Generating mesh visualization points")
            
            # Generate points for left arm meshes (red)
            for i, mesh in enumerate(left_arm_meshes):
                mesh_points = self._generate_mesh_surface_points(mesh, color=[255, 0, 0])
                if len(mesh_points) > 0:
                    mesh_vis_points = np.vstack([mesh_vis_points, mesh_points]) if len(mesh_vis_points) > 0 else mesh_points
            
            # Generate points for right arm meshes (blue)
            for i, mesh in enumerate(right_arm_meshes):
                mesh_points = self._generate_mesh_surface_points(mesh, color=[0, 0, 255])
                if len(mesh_points) > 0:
                    mesh_vis_points = np.vstack([mesh_vis_points, mesh_points]) if len(mesh_vis_points) > 0 else mesh_points
        
        # Use vectorized filtering for better performance
        keep_mask = self.filterByMeshes(points_xyz, left_arm_meshes, right_arm_meshes, buffer_m)
        
        filtered_pcd = pcd_array[keep_mask]
        
        # Add mesh visualization points if enabled
        if MESH_VISUALIZATION_ENABLED and len(mesh_vis_points) > 0:
            # Ensure data types match
            if filtered_pcd.dtype != mesh_vis_points.dtype:
                mesh_vis_points = mesh_vis_points.astype(filtered_pcd.dtype)
            
            # Combine filtered points with mesh visualization
            filtered_pcd = np.vstack([filtered_pcd, mesh_vis_points])
            
            if DEBUG_PRINT_MESH_INFO:
                print(f"DEBUG: Added {len(mesh_vis_points)} mesh visualization points")
        
        return filtered_pcd
    
    def trim_with_cylinder(self, pcd_array: np.ndarray, joint_states: np.ndarray, buffer_mm: float = 100.0) -> np.ndarray:
        """
        Trim points using cylindrical approximation of arm links with torus joints.
        
        Args:
            pcd_array: Point cloud array of shape (N, 6) with [x, y, z, r, g, b]
            joint_states: Array of shape (16,) containing [left_arm_7_joints, left_gripper_1, right_arm_7_joints, right_gripper_1]
            buffer_mm: Cylinder radius in millimeters
            
        Returns:
            Filtered point cloud array
        """
        if pcd_array.shape[0] == 0:
            return pcd_array
            
        points_xyz = pcd_array[:, :3]
        radius_m = buffer_mm / 1000.0  # Convert to meters
        
        # Split joint states for left and right arms (excluding gripper values)
        # joint_states format: [left_arm_7_joints, left_gripper_1, right_arm_7_joints, right_gripper_1]
        left_joint_states = joint_states[:7]  # First 7 values are left arm joints
        right_joint_states = joint_states[8:15]  # Values 8-14 are right arm joints (skip gripper at index 7)

        # ---------------- TIMING ----------------
        t0 = time.time()
        # Get arm spine points for both arms (now with 90-degree rotation applied)
        left_spine_points = self.spinePoints(left_joint_states, self.left_arm_base_pose)
        t1 = time.time()
        right_spine_points = self.spinePoints(right_joint_states, self.right_arm_base_pose)
        t2 = time.time()

        # Use vectorized filtering for better performance
        keep_mask = self.filterDualPoints(points_xyz, left_spine_points, right_spine_points, radius_m)
        t3 = time.time()

        print(f"TIMING -> spineLeft: {(t1 - t0):.4f}s  spineRight: {(t2 - t1):.4f}s  filter: {(t3 - t2):.4f}s  total: {(t3 - t0):.4f}s")
        
        return pcd_array[keep_mask]

    def filterDualPoints(self, points: np.ndarray, left_spine: np.ndarray, right_spine: np.ndarray, radius: float, batch_size: int = 500000) -> np.ndarray:
        """Fast point filtering against both arms using a fully-vectorized capsule test.

        The arm is approximated by a chain of capsules (cylinders with hemispherical caps).
        Each segment between consecutive spine points defines one capsule with the given
        radius.  A point is removed when it lies *inside* ANY capsule of either arm.

        Args:
            points: (N,3) array of XYZ positions.
            left_spine: (L,3) array of left-arm spine positions.
            right_spine: (R,3) array of right-arm spine positions.
            radius: capsule radius in metres.
            batch_size: number of points processed per batch to bound memory usage.

        Returns:
            Boolean numpy mask of length N where True means the point is kept.
        """
        global DEBUG_PRINT_FILTERING_INFO

        start_total = time.time()

        # Early exit if there are no spine segments to test against
        if (left_spine is None or len(left_spine) < 2) and (right_spine is None or len(right_spine) < 2):
            return np.ones(len(points), dtype=bool)

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        radius2 = radius * radius  # compare squared distances to avoid sqrt

        # Helper: build segment start/end tensors of shape (S,3)
        def _segments_from_spine(spine: np.ndarray) -> torch.Tensor:
            if spine is None or len(spine) < 2:
                return torch.empty((0, 3), dtype=torch.float32, device=device), torch.empty((0, 3), dtype=torch.float32, device=device)
            starts = torch.from_numpy(spine[:-1]).float().to(device)
            ends = torch.from_numpy(spine[1:]).float().to(device)
            return starts, ends

        l_start, l_end = _segments_from_spine(left_spine)
        r_start, r_end = _segments_from_spine(right_spine)

        # Concatenate segments from both arms => (S,3)
        seg_starts = torch.cat([l_start, r_start], dim=0)
        seg_ends = torch.cat([l_end, r_end], dim=0)
        seg_count = seg_starts.shape[0]

        if seg_count == 0:
            return np.ones(len(points), dtype=bool)

        # Pre-compute segment data
        seg_vec = seg_ends - seg_starts                      # (S,3)
        seg_len = torch.norm(seg_vec, dim=1, keepdim=True)   # (S,1)
        seg_unit = seg_vec / (seg_len + 1e-9)                # (S,3)

        if DEBUG_PRINT_FILTERING_INFO:
            print(f"DEBUG: Vectorised capsule test on {len(points)} points, {seg_count} segments, radius={radius}")
            print(f"DEBUG: Processing in batches of {batch_size} on device {device}")

        # Allocate output mask on CPU for easy numpy interoperability
        keep_mask_cpu = np.ones(len(points), dtype=bool)

        # Convert points to tensor once, then slice per batch
        points_tensor = torch.from_numpy(points).float().to(device)

        with torch.no_grad():
            for start_idx in range(0, len(points), batch_size):
                end_idx = min(start_idx + batch_size, len(points))
                batch = points_tensor[start_idx:end_idx]          # (B,3)

                batch_start_t = time.time()
                # (B,1,3) - (1,S,3) => (B,S,3)
                diff = batch.unsqueeze(1) - seg_starts.unsqueeze(0)

                # Project diff onto segment unit
                proj = (diff * seg_unit.unsqueeze(0)).sum(dim=2, keepdim=True)  # (B,S,1)
                # First clamp to the lower bound (0), then apply the upper bound (segment length)
                proj_clamped = torch.clamp(proj, min=0.0)
                proj_clamped = torch.clamp(proj_clamped, max=seg_len.unsqueeze(0))

                # Closest points on segment to each point
                closest = seg_starts.unsqueeze(0) + proj_clamped * seg_unit.unsqueeze(0)  # (B,S,3)

                # Squared distances to closest points
                dist2 = torch.sum((batch.unsqueeze(1) - closest) ** 2, dim=2)  # (B,S)

                # Minimum distance squared over all capsules
                min_dist2, _ = torch.min(dist2, dim=1)  # (B,)

                # Keep if outside all capsules
                keep_batch = min_dist2 > radius2

                # Write back to CPU mask
                keep_mask_cpu[start_idx:end_idx] = keep_batch.cpu().numpy()

                batch_end_t = time.time()
                if DEBUG_PRINT_FILTERING_INFO:
                    print(f"    Batch {start_idx}-{end_idx} took {(batch_end_t - batch_start_t):.4f}s")

        if DEBUG_PRINT_FILTERING_INFO:
            removed = len(points) - np.count_nonzero(keep_mask_cpu)
            print(f"DEBUG: Capsule filtering removed {removed} / {len(points)} points")
            print(f"DEBUG: Total capsule filtering time {(time.time() - start_total):.4f}s")

        return keep_mask_cpu
    
    def genDualSurfacePoints(self, joint_states: np.ndarray, num_points: int = 5000, radius: float = 0.05, reference_pcd: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Generate surface points on both arms for visualization with torus joints.
        
        Args:
            joint_states: Array of shape (16,) containing [left_arm_7_joints, left_gripper_1, right_arm_7_joints, right_gripper_1]
            num_points: Total number of points to generate
            reference_pcd: Reference point cloud to match color format and data type
            
        Returns:
            Array of surface points with shape (N, 6) [x, y, z, r, g, b]
        """
        points_per_arm = num_points // 2
        
        # Determine color format based on reference point cloud
        if reference_pcd is not None and reference_pcd.shape[0] > 0:
  
            print(f"trimArmPcd.py: genDualSurfacePoints - Using 0-255 colors, scalePcd will handle scaling")
        else:

            print(f"trimArmPcd.py: genDualSurfacePoints - No reference PCD, using 0-255 colors")
                    # Default to 0-255 range
        left_color = [255, 0, 0]  # Red
        right_color = [0, 0, 255]  # Blue
        # Split joint states for left and right arms (excluding gripper values)
        left_joint_states = joint_states[:7]  # First 7 values are left arm joints
        right_joint_states = joint_states[8:15]  # Values 8-14 are right arm joints (skip gripper at index 7)
        
        # Generate points for left arm (red color) - now with 90-degree rotation
        left_spine_points = self.spinePoints(left_joint_states, self.left_arm_base_pose)
        left_surface_points = self.genSurfacePoints(
            left_spine_points, points_per_arm, color=left_color, radius=radius
        )
        
        # Generate points for right arm (blue color) - now with 90-degree rotation
        right_spine_points = self.spinePoints(right_joint_states, self.right_arm_base_pose)
        right_surface_points = self.genSurfacePoints(
            right_spine_points, points_per_arm, color=right_color, radius=radius
        )
        
        # Combine both arms
        if len(left_surface_points) > 0 and len(right_surface_points) > 0:
            combined_points = np.vstack([left_surface_points, right_surface_points])
        elif len(left_surface_points) > 0:
            combined_points = left_surface_points
        elif len(right_surface_points) > 0:
            combined_points = right_surface_points
        else:
            combined_points = np.empty((0, 6))
        
        # Match data type of reference point cloud if provided
        if reference_pcd is not None and reference_pcd.shape[0] > 0 and combined_points.shape[0] > 0:
            combined_points = combined_points.astype(reference_pcd.dtype)
        
        # 统一使用 scalePcd 处理颜色缩放
        combined_points = scalePcd(combined_points, tag="trimArmPcd_surface_points")
        
        return combined_points
    
    def _get_arm_meshes(self, joint_states: np.ndarray, base_pose: np.ndarray, is_left_arm: bool) -> List[trimesh.Trimesh]:
        """Get mesh geometries for arm links at given joint configuration."""
        meshes = []
        
        # Use mesh-enabled robot if available
        robot_to_use = self.robot_with_meshes if self.robot_with_meshes is not None else self.robot
        
        try:
            # Determine which arm transform to use
            arm_transform = self.left_arm_transform if is_left_arm else self.right_arm_transform
            
            # Create joint configuration dictionary
            joint_config = {}
            for i, joint_name in enumerate(self.arm_joint_names):
                joint_config[joint_name] = joint_states[i]
            
            # Update robot configuration
            robot_to_use.update_cfg(joint_config)
            
            # Get mesh for each arm link
            for link_name in self.arm_link_names:
                try:
                    # Get the link transform from base frame
                    link_transform = robot_to_use.get_transform(frame_to=link_name)
                    
                    # Apply arm transformation (includes rotation and translation)
                    final_transform = arm_transform @ link_transform
                    
                    # Get the link object
                    link = robot_to_use.link_map.get(link_name)
                    if link is None:
                        continue
                    
                    # Process visual geometries
                    for visual in link.visuals:
                        if visual.geometry.mesh is not None:
                            # Try to get mesh from geometry
                            mesh = self._extract_mesh_from_geometry(visual.geometry.mesh, final_transform)
                            if mesh is not None:
                                meshes.append(mesh)
                        else:
                            # Handle primitive geometries (box, cylinder, sphere)
                            mesh = self._create_primitive_mesh(visual.geometry, final_transform)
                            if mesh is not None:
                                meshes.append(mesh)
                    
                    # Process collision geometries if no visual geometries found
                    if not meshes and link.collisions:
                        for collision in link.collisions:
                            if collision.geometry.mesh is not None:
                                mesh = self._extract_mesh_from_geometry(collision.geometry.mesh, final_transform)
                                if mesh is not None:
                                    meshes.append(mesh)
                            else:
                                # Handle primitive collision geometries
                                mesh = self._create_primitive_mesh(collision.geometry, final_transform)
                                if mesh is not None:
                                    meshes.append(mesh)
                
                except Exception as e:
                    logger.debug(f"Could not process link {link_name}: {e}")
                    continue
        
        except Exception as e:
            logger.warning(f"Could not load arm meshes: {e}")
            # Fallback: return empty list, will use cylindrical approximation
            return []
        
        logger.info(f"Loaded {len(meshes)} meshes for {'left' if is_left_arm else 'right'} arm")
        return meshes
    
    def _extract_mesh_from_geometry(self, mesh_geometry, transform: np.ndarray) -> Optional[trimesh.Trimesh]:
        """Extract mesh from geometry object."""
        global MESH_SIMPLIFY_TARGET_FACES, MESH_SIMPLIFY_MAX_FACES, DEBUG_PRINT_MESH_INFO
        
        try:
            # Try to get mesh from filename
            if hasattr(mesh_geometry, 'filename') and mesh_geometry.filename:
                mesh_path = mesh_geometry.filename
                
                # If path is relative, make it relative to the URDF directory
                if not os.path.isabs(mesh_path):
                    urdf_dir = os.path.dirname(self.urdf_path)
                    mesh_path = os.path.join(urdf_dir, mesh_path)
                
                if os.path.exists(mesh_path):
                    mesh = trimesh.load(mesh_path, force='mesh')
                    if isinstance(mesh, trimesh.Trimesh):
                        original_faces = len(mesh.faces)
                        if DEBUG_PRINT_MESH_INFO:
                            print(f"DEBUG: Loaded mesh from {mesh_path} with {original_faces} faces")
                        
                        # HEAVILY SIMPLIFY MESH FOR SPEED - use global variables
                        bbox_size = mesh.bounding_box.extents
                        max_dimension = max(bbox_size)
                        
                        # Use global target faces setting
                        target_faces = max(MESH_SIMPLIFY_TARGET_FACES, int((max_dimension / 0.1) ** 1.5))
                        target_faces = min(target_faces, MESH_SIMPLIFY_MAX_FACES)
                        
                        if len(mesh.faces) > target_faces:
                            # Use aggressive mesh simplification
                            try:
                                # Use correct trimesh simplification method
                                simplified = mesh.simplify_quadric_decimation(target_faces)
                                if simplified is not None:
                                    mesh = simplified
                                    if DEBUG_PRINT_MESH_INFO:
                                        print(f"DEBUG: Simplified mesh to {len(mesh.faces)} faces (target: {target_faces})")
                                else:
                                    raise Exception("Simplification returned None")
                            except:
                                # Fallback: use convex hull as ultimate fallback
                                try:
                                    mesh = mesh.convex_hull
                                    if DEBUG_PRINT_MESH_INFO:
                                        print(f"DEBUG: Used convex hull fallback with {len(mesh.faces)} faces")
                                except:
                                    if DEBUG_PRINT_MESH_INFO:
                                        print("DEBUG: Mesh simplification failed, using original mesh")
                        
                        # Apply scale if specified
                        if hasattr(mesh_geometry, 'scale') and mesh_geometry.scale is not None:
                            mesh.apply_scale(mesh_geometry.scale)
                        # Apply transformations
                        mesh.apply_transform(transform)
                        return mesh
            
            # Try to get mesh directly from geometry if it's already loaded
            if hasattr(mesh_geometry, 'mesh') and mesh_geometry.mesh is not None:
                mesh = mesh_geometry.mesh
                if isinstance(mesh, trimesh.Trimesh):
                    mesh = mesh.copy()  # Make a copy to avoid modifying original
                    original_faces = len(mesh.faces)
                    if DEBUG_PRINT_MESH_INFO:
                        print(f"DEBUG: Using direct mesh with {original_faces} faces")
                    
                    # HEAVILY SIMPLIFY MESH FOR SPEED - use global variables
                    bbox_size = mesh.bounding_box.extents
                    max_dimension = max(bbox_size)
                    target_faces = max(MESH_SIMPLIFY_TARGET_FACES, int((max_dimension / 0.1) ** 1.5))
                    target_faces = min(target_faces, MESH_SIMPLIFY_MAX_FACES)
                    
                    if len(mesh.faces) > target_faces:
                        try:
                            mesh = mesh.simplify_quadric_decimation(target_faces)
                            if DEBUG_PRINT_MESH_INFO:
                                print(f"DEBUG: Simplified direct mesh to {len(mesh.faces)} faces (target: {target_faces})")
                        except:
                            try:
                                mesh = mesh.convex_hull
                                if DEBUG_PRINT_MESH_INFO:
                                    print(f"DEBUG: Used convex hull fallback with {len(mesh.faces)} faces")
                            except:
                                if DEBUG_PRINT_MESH_INFO:
                                    print("DEBUG: Direct mesh simplification failed, using original")
                    
                    # Apply scale if specified
                    if hasattr(mesh_geometry, 'scale') and mesh_geometry.scale is not None:
                        mesh.apply_scale(mesh_geometry.scale)
                    # Apply transformations
                    mesh.apply_transform(transform)
                    return mesh
            
            return None
            
        except Exception as e:
            if DEBUG_PRINT_MESH_INFO:
                logger.debug(f"Could not extract mesh from geometry: {e}")
            return None
    
    def _create_primitive_mesh(self, geometry, transform: np.ndarray) -> Optional[trimesh.Trimesh]:
        """Create mesh from primitive geometry (box, cylinder, sphere)."""
        try:
            if hasattr(geometry, 'box') and geometry.box is not None:
                # Box geometry
                size = geometry.box.size
                mesh = trimesh.creation.box(extents=size)
            elif hasattr(geometry, 'cylinder') and geometry.cylinder is not None:
                # Cylinder geometry
                radius = geometry.cylinder.radius
                length = geometry.cylinder.length
                mesh = trimesh.creation.cylinder(radius=radius, height=length)
            elif hasattr(geometry, 'sphere') and geometry.sphere is not None:
                # Sphere geometry
                radius = geometry.sphere.radius
                mesh = trimesh.creation.icosphere(radius=radius)
            else:
                return None
            
            # Apply transform
            mesh.apply_transform(transform)
            return mesh
            
        except Exception as e:
            logger.debug(f"Could not create primitive mesh: {e}")
            return None
    
    def _generate_mesh_surface_points(self, mesh: trimesh.Trimesh, num_points_per_face: Optional[int] = None, 
                                    color: Optional[Sequence[Union[int, float]]] = None) -> np.ndarray:
        """Generate surface points from mesh for visualization."""
        global MESH_SURFACE_POINTS_PER_FACE, DEBUG_PRINT_MESH_INFO
        
        if num_points_per_face is None:
            num_points_per_face = MESH_SURFACE_POINTS_PER_FACE
        
        if color is None:
            color = [0, 255, 0]  # Default green color
        
        try:
            # Calculate total number of points needed
            num_faces = len(mesh.faces)
            total_points = num_faces * num_points_per_face
            
            if DEBUG_PRINT_MESH_INFO:
                print(f"DEBUG: Generating {total_points} surface points from {num_faces} faces")
            
            # Sample points on mesh surface
            surface_points, face_indices = mesh.sample(total_points, return_index=True)
            
            # Add colors to the points
            points_with_color = np.column_stack([
                surface_points,
                np.tile(color, (len(surface_points), 1))
            ])
            
            return points_with_color.astype(np.float32)
            
        except Exception as e:
            if DEBUG_PRINT_MESH_INFO:
                logger.warning(f"Could not generate mesh surface points: {e}")
            return np.empty((0, 6), dtype=np.float32)
    
    def filterByMeshes(self, points: np.ndarray, left_meshes: List[trimesh.Trimesh], 
                                   right_meshes: List[trimesh.Trimesh], buffer: float) -> np.ndarray:
        """Filter points using mesh collision detection with PyTorch acceleration."""
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Convert points to PyTorch tensor
        points_tensor = torch.from_numpy(points).float().to(device)
        keep_mask = torch.ones(len(points), dtype=torch.bool, device=device)
        
        # Process left arm meshes
        for mesh in left_meshes:
            keep_mask &= self.filterByMesh(points_tensor, mesh, buffer, device)
        
        # Process right arm meshes
        for mesh in right_meshes:
            keep_mask &= self.filterByMesh(points_tensor, mesh, buffer, device)
        
        return keep_mask.cpu().numpy()
    
    def filterByMesh(self, points: torch.Tensor, mesh: trimesh.Trimesh, 
                                          buffer: float, device: torch.device) -> torch.Tensor:
        """Filter points against a single mesh using PyTorch."""
        try:
            # Convert mesh to PyTorch tensors
            vertices = torch.from_numpy(mesh.vertices).float().to(device)
            faces = torch.from_numpy(mesh.faces).long().to(device)
            
            # Use trimesh's built-in distance calculation (fallback to CPU)
            points_cpu = points.cpu().numpy()
            distances = mesh.nearest.signed_distance(points_cpu)
            distances_tensor = torch.from_numpy(distances).float().to(device)
            
            # Points are kept if they are outside the mesh + buffer
            keep_mask = distances_tensor > buffer
            
            return keep_mask
            
        except Exception as e:
            logger.debug(f"Error in mesh filtering: {e}")
            # Return all points as kept if error occurs
            return torch.ones(len(points), dtype=torch.bool, device=device)

    def spinePoints(self, joint_states: np.ndarray, base_pose: np.ndarray) -> np.ndarray:
        """Get spine points (joint positions) for cylindrical approximation."""
        global DEBUG_PRINT_SPINE_INFO
        
        spine_points = []
        
        # Determine which arm this is based on base pose
        is_left_arm = np.allclose(base_pose, self.left_arm_base_pose)
        arm_transform = self.left_arm_transform if is_left_arm else self.right_arm_transform
        
        if DEBUG_PRINT_SPINE_INFO:
            print(f"DEBUG: Computing spine points for {'left' if is_left_arm else 'right'} arm")
            print(f"DEBUG: Base pose: {base_pose}")
            print(f"DEBUG: Joint states: {joint_states}")
        
        # Create joint configuration dictionary
        joint_config = {}
        for i, joint_name in enumerate(self.arm_joint_names):
            joint_config[joint_name] = joint_states[i]
        
        # Update robot configuration
        try:
            self.robot.update_cfg(joint_config)
        except Exception as e:
            logger.warning(f"Could not update robot configuration: {e}")
            # Fallback to simple spine approximation with rotation
            spine_points.append(base_pose)
            # Add some default points along a simple arm configuration
            for i in range(7):
                # Create point in local arm coordinate system
                local_point = np.array([0, 0, i * 0.1, 1])  # Homogeneous coordinates
                # Transform to world coordinates with rotation
                world_point = arm_transform @ local_point
                spine_points.append(world_point[:3])
            
            if DEBUG_PRINT_SPINE_INFO:
                print(f"DEBUG: Using fallback spine with {len(spine_points)} points")
                for i, point in enumerate(spine_points):
                    print(f"DEBUG: Spine point {i}: {point}")
            
            return np.array(spine_points)
        
        # Add base position
        spine_points.append(base_pose)
        
        # Get position of each joint/link
        for link_name in ['panda_link1', 'panda_link2', 'panda_link3', 'panda_link4', 
                         'panda_link5', 'panda_link6', 'panda_link7', 'panda_hand']:
            try:
                # Get transform from base frame to this link
                link_pose = self.robot.get_transform(frame_to=link_name)
                
                # Apply arm transformation (includes rotation and translation)
                final_pose = arm_transform @ link_pose
                
                # Extract position
                position = final_pose[:3, 3]
                spine_points.append(position)
                
            except Exception as e:
                print(traceback.format_exc())
                logger.warning(f"Could not get transform for link {link_name}: {e}")
                continue
        
        if DEBUG_PRINT_SPINE_INFO:
            print(f"DEBUG: Generated {len(spine_points)} spine points")
            for i, point in enumerate(spine_points):
                print(f"DEBUG: Spine point {i}: {point}")
        
        return np.array(spine_points)
    
    def genSurfacePoints(self, spine_points: np.ndarray, num_points: int, color: Sequence[Union[int, float]], radius: float) -> np.ndarray:
        """Generate surface points around the arm spine with bullet shapes."""
        global spine_vis
        
        if len(spine_points) < 2:
            return np.empty((0, 6))
        
        surface_points = []
        
        # Only generate points if spine visualization is enabled
        if not spine_vis:
            return np.empty((0, 6))
        
        # Calculate points per segment
        num_segments = len(spine_points) - 1
        points_per_segment = max(1, num_points // num_segments)
        
        # Generate points for bullet-shaped segments
        for i in range(num_segments):
            p1 = spine_points[i]
            p2 = spine_points[i + 1]
            
            # Determine if this is the last segment
            is_last_segment = (i == num_segments - 1)
            
            # Generate points around this bullet segment
            segment_points = self.genBulletSurface(p1, p2, points_per_segment, color, radius, is_last_segment)
            surface_points.extend(segment_points)
        
        return np.array(surface_points) if surface_points else np.empty((0, 6))
    
    def genBulletSurface(self, p1: np.ndarray, p2: np.ndarray, num_points: int, 
                        color: Sequence[Union[int, float]], radius: float, is_last_segment: bool = False) -> List[np.ndarray]:
        """Generate surface points on bullet shape (cylinder + hemisphere top)."""
        surface_points = []
        
        # Calculate segment vector
        segment_vec = p2 - p1
        segment_length = np.linalg.norm(segment_vec)
        
        if segment_length < 1e-6:
            return surface_points
        
        segment_unit = segment_vec / segment_length
        
        # Create perpendicular vectors for cylinder surface
        if abs(segment_unit[2]) < 0.9:
            perp1 = np.cross(segment_unit, [0, 0, 1])
        else:
            perp1 = np.cross(segment_unit, [1, 0, 0])
        perp1 = perp1 / np.linalg.norm(perp1)
        perp2 = np.cross(segment_unit, perp1)
        
        # Split points between cylinder and hemisphere
        if is_last_segment:
            # Last segment: only cylinder, no hemisphere
            cylinder_points = num_points
            hemisphere_points = 0
        else:
            # Regular segment: cylinder + hemisphere
            cylinder_points = num_points // 2
            hemisphere_points = num_points - cylinder_points
        
        # Generate cylindrical surface points
        for j in range(cylinder_points):
            # Random position along segment
            t = np.random.uniform(0, 1)
            center = p1 + t * segment_vec
            
            # Random angle around cylinder
            angle = np.random.uniform(0, 2 * np.pi)
            
            # Point on cylinder surface
            surface_point = center + radius * (np.cos(angle) * perp1 + np.sin(angle) * perp2)
            
            # Add color
            point_with_color = np.concatenate([surface_point, color])
            surface_points.append(point_with_color)
        
        # Generate hemispherical surface points (unless last segment)
        if not is_last_segment:
            for j in range(hemisphere_points):
                # Random point on hemisphere using spherical coordinates
                # phi: angle from north pole (0 to pi/2 for hemisphere)
                phi = np.random.uniform(0, np.pi / 2)
                # theta: azimuthal angle (0 to 2*pi)
                theta = np.random.uniform(0, 2 * np.pi)
                
                # Convert to Cartesian coordinates (hemisphere centered at p2)
                local_x = radius * np.sin(phi) * np.cos(theta)
                local_y = radius * np.sin(phi) * np.sin(theta)
                local_z = radius * np.cos(phi)
                
                # Transform to world coordinates
                # local_z is along segment_unit direction (upward)
                # local_x and local_y are in the perpendicular plane
                surface_point = (p2 + 
                               local_x * perp1 + 
                               local_y * perp2 + 
                               local_z * segment_unit)
                
                # Add color
                point_with_color = np.concatenate([surface_point, color])
                surface_points.append(point_with_color)
        
        return surface_points


@timeit(color='magenta')
def trimArmPcd(
    unifiedData: Dict[str, np.ndarray], 
    param: str,
    urdf_path: str,
    pcd_keys: List[str]
) -> Dict[str, np.ndarray]:
    """
    Trim arm points from point cloud data based on URDF and joint states.
    
    Args:
        unifiedData: Dictionary containing single frame data with keys:
                    ["rgb":np.ndarray,"depth":np.ndarray,"pointcloud":np.ndarray,
                     "segpointcloud":np.ndarray,"state":np.ndarray,"action":np.ndarray,
                     "endpose":np.ndarray,"cameras":dict,"gaussian":dict]
                    where pointcloud and segpointcloud are shape (Np, 6) for single frame
        param: Trimming parameter string:
               - 'm100': Use mesh-based trimming with 100mm buffer
               - 'c100': Use cylindrical trimming with 100mm buffer(radius)
        urdf_path: Path to robot URDF file
        pcd_keys: List of point cloud keys to process
        
    Returns:
        Dictionary containing original data with trimmed point clouds
    """
   
    if param.startswith('c') or param.startswith('m'):
        buffer = float(param[1:]) 
    else:
        logger.error(f"Invalid parameter format: {param}. Use 'm100' for mesh or 'c100' for cylindrical")
        return unifiedData
    
    if 'state' in unifiedData and unifiedData['state'] is not None:
        joint_states = unifiedData['state']
    else:
        logger.error("No joint states provided and cannot extract from unifiedData")
        return unifiedData
    
    # Validate joint states shape
    if len(joint_states) != 16:
        logger.error(f"Expected 16 joint states (7 left + 7 right + 2 grippers), got {len(joint_states)}")
        return unifiedData
    
    # Initialize trimmer (global initialization without joint_states)
    trimmer = ArmPointCloudTrimmer(urdf_path)
    
    # Process each point cloud key
    for pcd_key in pcd_keys:
        if pcd_key in unifiedData.keys():
            point_cloud_data = unifiedData[pcd_key]
            if point_cloud_data is not None:
                # Validate input shape
                if len(point_cloud_data.shape) != 2 or point_cloud_data.shape[1] != 6:
                    logger.warning(f"Skipping {pcd_key}: Expected shape (Np, 6), got {point_cloud_data.shape}")
                    continue
                
                original_count = point_cloud_data.shape[0]
                
                # Apply appropriate trimming method (now passing joint_states)
                if param[0] == 'm':
                    unifiedData[pcd_key] = trimmer.trim_with_mesh(point_cloud_data, joint_states, buffer)
                elif param[0] == 'c':
                    unifiedData[pcd_key] = trimmer.trim_with_cylinder(point_cloud_data, joint_states, buffer)
                trimmed_count = unifiedData[pcd_key].shape[0]
                
                # If trimmed to no points, add a point at (0,0,0)
                if trimmed_count == 0:
                    # Create a point at origin with default color (white: 255,255,255)
                    origin_point = np.array([[0.0, 0.0, 0.0, 255, 255, 255]], dtype=point_cloud_data.dtype)
                    unifiedData[pcd_key] = origin_point
                    trimmed_count = 1
                    print(f"    ⚠️  Added origin point (0,0,0) to {pcd_key} as fallback")
                
                print(f"    🤖 Arm trimmed {pcd_key}: {original_count} → {trimmed_count} points using {param[0]} method")
                
                # Generate surface points
                surface_points = trimmer.genDualSurfacePoints(joint_states, 100000, radius=buffer/1000.0, reference_pcd=point_cloud_data)
                
                # Ensure both arrays have the same data type before concatenation
                if unifiedData[pcd_key].dtype != surface_points.dtype:
                    surface_points = surface_points.astype(unifiedData[pcd_key].dtype)
                
                unifiedData[pcd_key]=np.vstack([unifiedData[pcd_key], surface_points])#LEAVE IT, NO DELETE!!!!!!!
    
    return unifiedData


if __name__ == "__main__":
    import time
    
    # Create some dummy test data
    print("Creating test data...")
    
    # Create dummy point cloud data
    np.random.seed(42)
    num_points = 10000
    dummy_pointcloud = np.random.randn(num_points, 6).astype(np.float32)
    dummy_pointcloud[:, :3] *= 0.5  # Scale positions to reasonable range
    dummy_pointcloud[:, 3:] = np.random.randint(0, 256, (num_points, 3))  # RGB colors
    
    # Create dummy joint states (16 values: 7 left + 1 gripper + 7 right + 1 gripper)
    dummy_joint_states = np.random.uniform(-1.0, 1.0, 16).astype(np.float32)
    
    # Create unified data structure
    unifiedData = {
        'pointcloud': dummy_pointcloud,
        'segpointcloud': dummy_pointcloud.copy(),
        'state': dummy_joint_states,
        'rgb': np.random.randint(0, 256, (84, 84, 3), dtype=np.uint8),
        'depth': np.random.uniform(0.1, 2.0, (84, 84)).astype(np.float32)
    }
    
    # Test parameters
    urdf_path = str(Path(__file__).parent.parent / "RoboTwin_Close_old/assets/embodiments/franka-panda/panda.urdf")
    param = "c150"  # Cylindrical trimming with 150mm radius
    pcd_keys = ['pointcloud', 'segpointcloud']
    
    print(f"Testing trimArmPcd with {num_points} points...")
    print(f"URDF path: {urdf_path}")
    print(f"Parameter: {param}")
    print(f"Point cloud keys: {pcd_keys}")
    
    # Check if URDF file exists
    if not os.path.exists(urdf_path):
        print(f"ERROR: URDF file not found at {urdf_path}")
        print("Please provide a valid URDF path")
        exit(1)
    
    try:
        # Run the trimming function
        start_time = time.time()
        result = trimArmPcd(unifiedData, param, urdf_path, pcd_keys)
        total_time = time.time() - start_time
        
        print(f"\nTotal execution time: {total_time:.3f}s")
        print("\nResults:")
        for key in pcd_keys:
            if key in result:
                original_count = unifiedData[key].shape[0]
                final_count = result[key].shape[0]
                print(f"  {key}: {original_count} -> {final_count} points")
        
        print("\nTest completed successfully!")
        
    except Exception as e:
        print(f"ERROR during testing: {e}")
        import traceback
        traceback.print_exc()

